import math
import numpy as np
# import open3d as o3d


def o3d_read_mesh(mesh_path, scale=None):
    mesh = o3d.io.read_triangle_mesh(mesh_path)
    mesh.compute_vertex_normals()
    if scale is not None:
        mesh.scale(scale, center=[0,0,0])
    return mesh

def get_4x4_matrix(quat, pos):
    '''
    quat: 四元数, q = [x, y, z, w]
    pos: p = [x, y, z]
    '''
    if isinstance(quat, list):
        quat = np.array(quat, np.float64)
    if isinstance(pos, list):
        pos = np.array(pos, np.float64)
    t = np.eye(4).astype(np.float64)
    t[:3, 3] = pos
    t[:3, :3] = q2R(quat)
    return t

def q2R(q):
    """
    q: 四元数, q = [x, y, z, w]
    R: 旋转矩阵,左乘
    """
    q = q[[3,0,1,2]]  # (x,y,z,w) -> (w,x,y,z)
    R = np.zeros((3, 3), dtype=float)

    R[0, 0] = 1 - 2 * (q[2] ** 2 + q[3] ** 2)
    R[1, 1] = 1 - 2 * (q[1] ** 2 + q[3] ** 2)
    R[2, 2] = 1 - 2 * (q[1] ** 2 + q[2] ** 2)

    R[0, 1] = 2 * (q[1] * q[2] - q[0] * q[3])
    R[1, 0] = 2 * (q[1] * q[2] + q[0] * q[3])

    R[0, 2] = 2 * (q[1] * q[3] + q[0] * q[2])
    R[2, 0] = 2 * (q[1] * q[3] - q[0] * q[2])

    R[1, 2] = 2 * (q[2] * q[3] - q[0] * q[1])
    R[2, 1] = 2 * (q[2] * q[3] + q[0] * q[1])

    return R

def R2q(R:np.array):
    """
    R: 旋转矩阵,左乘
    q: 四元数, q = [x, y, z, w]
    """
        
    w = R[0,0]+R[1,1]+R[2,2]+1
    x = R[0,0]-R[1,1]-R[2,2]+1
    y = -R[0,0]+R[1,1]-R[2,2]+1
    z = -R[0,0]-R[1,1]+R[2,2]+1
    
    q = np.array([w,x,y,z])
    index = np.argmax(q)
    # assert q[index]>0, "max(q) > 0"
    if q[index]<=0:
        return np.array([0,0,0,1],dtype=np.float32)

    q[index] = math.sqrt(q[index]) / 2

    if index==0:
        q0 = q[index]
        q1 = (R[2,1]-R[1,2]) / (4*q0)
        q2 = (R[0,2]-R[2,0]) / (4*q0)
        q3 = (R[1,0]-R[0,1]) / (4*q0)
    elif index==1:
        q1 = q[index]
        q0 = (R[2,1]-R[1,2]) / (4*q1)
        q2 = (R[0,1]+R[1,0]) / (4*q1)
        q3 = (R[2,0]+R[0,2]) / (4*q1)
    elif index==2:
        q2 = q[index]
        q0 = (R[0,2]-R[2,0]) / (4*q2)
        q1 = (R[0,1]+R[1,0]) / (4*q2)
        q3 = (R[1,2]+R[2,1]) / (4*q2)
    elif index==3:
        q3 = q[index]
        q0 = (R[1,0]-R[0,1]) / (4*q3)
        q1 = (R[2,0]+R[0,2]) / (4*q3)
        q2 = (R[1,2]+R[2,1]) / (4*q3)
    else:
        raise ValueError('index error:' +  str(index))
    return np.array([q1,q2,q3,q0],dtype=np.float32)

def transform_coordinates_3d(coordinates, sRT):
    """
    Args:
        coordinates: [3, N]
        sRT: [4, 4]

    Returns:
        new_coordinates: [3, N]

    """
    assert coordinates.shape[0] == 3
    coordinates = np.vstack([coordinates, np.ones((1, coordinates.shape[1]), dtype=np.float32)])
    new_coordinates = sRT @ coordinates
    new_coordinates = new_coordinates[:3, :] / new_coordinates[3, :]
    return new_coordinates

# def draw_3Daxes(RT, size=0.01):
#     '''
#     Args:
#         RT: [4x4] numpy
#         size: length for axis
#     Return:
#         open3D lineset object
#     '''
#     xyz_axis = np.array([[0, 0, 0], [size, 0, 0], [0, size, 0], [0, 0, size]], dtype=np.float32).transpose()
#     RT_xyz_axis = transform_coordinates_3d(xyz_axis, RT)
#     lineset = o3d.geometry.LineSet()
#     lineset.points = o3d.utility.Vector3dVector(RT_xyz_axis.transpose())
#     lineset.lines  = o3d.utility.Vector2iVector([[0, 1], [0, 2], [0, 3]])
#     lineset.colors = o3d.utility.Vector3dVector([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
#     return lineset